import DobotDllType as dType
import time
import numpy as np
from transform_matrix import matrix
from serial_receive import SerialReceiver
CON_STR = {
    dType.DobotConnect.DobotConnect_NoError: "DobotConnect_NoError",
    dType.DobotConnect.DobotConnect_NotFound: "DobotConnect_NotFound",
    dType.DobotConnect.DobotConnect_Occupied: "DobotConnect_Occupied"}


class ArmController:
    def __init__(self, port="COM9", baudrate=115200):
        # Load Dll and get the CDLL object
        self.api = dType.load()
        # Connect Dobot
        state = dType.ConnectDobot(self.api, port, baudrate)[0]
        print("Connect status:", CON_STR[state])
        if state != dType.DobotConnect.DobotConnect_NoError:
            raise Exception("Failed to connect to Dobot")
        dType.SetQueuedCmdClear(self.api)
        self.receiver = SerialReceiver(port="COM10",baudrate=9600,encoding="utf-8",group_size=4)
        self.receiver.find_port()
        self.receiver.start()

    def grasp_place(self, start_pos, target_pos, jump_height=20, velocity=50, acceleration=50):
        """
        使用JUMP模式从起始点移动到目标点
        
        参数:
            start_pos: 起始位置 [x, y, z, rHead]
            target_pos: 目标位置 [x, y, z, rHead]
            jump_height: 跳跃高度(mm)
            velocity: 速度百分比(0-100)
            acceleration: 加速度百分比(0-100)
        """
        dType.SetQueuedCmdClear(self.api)
        # 设置JUMP参数
        print(f"设置JUMP参数: 跳跃高度={jump_height}mm")
        dType.SetPTPJumpParams(self.api, jump_height, 400)

        # 设置运动速度和加速度
        print(f"设置运动速度: {velocity}%，加速度: {acceleration}%")
        dType.SetPTPCommonParams(self.api, velocity, acceleration)

        # 定义起始点和目标点
        start_x, start_y, start_z, start_r = start_pos
        target_x, target_y, target_z, target_r = target_pos
        # if (start_x>320 or start_y>320 or target_x>320 or target_y>320):
        #     print("起始点或目标点位置超出范围，跳过JUMP运动")
        #     return
        # if (start_x**2 + start_y**2) < 10000:
        #     print("起始点位置过近，跳过JUMP运动")
        #     return
        # if (target_x**2 + target_y**2) < 10000:
        #     print("目标点位置过近，跳过JUMP运动")
        #     return

        # 构建JUMP运动路径
        print("构建JUMP运动路径...")

        # 先移动到起始点(使用关节运动模式)
        print(f"移动到起始点: X={start_x}, Y={start_y}, Z={start_z}, R={start_r}")
        dType.SetPTPCmd(self.api, dType.PTPMode.PTPJUMPXYZMode, start_x, start_y, start_z, start_r, isQueued=1)
        # z=self.get_pose()[2]-self.receiver.get_current_distance()
        # dType.SetPTPCmd(self.api, dType.PTPMode.PTPMOVJXYZMode, start_x, start_y, z, start_r, isQueued=1)
        # dType.SetEndEffectorGripperEx(self.api, 1, 0, isQueued=1)  # 打开夹爪
        dType.SetEndEffectorGripperEx(self.api, 1, 1, isQueued=1)  # 关闭夹爪
        time.sleep(1)
        # dType.SetEndEffectorSuctionCup(self.api, 1, 1, isQueued=1)  # 打开吸盘
        # time.sleep(3)  # 等待吸盘关闭
        last = dType.SetPTPCmd(self.api, dType.PTPMode.PTPJUMPXYZMode, target_x, target_y, target_z, target_r, isQueued=1)[0]
        dType.SetEndEffectorGripperEx(self.api, 1, 0, isQueued=1)  # 打开夹爪
        # dType.SetEndEffectorSuctionCup(self.api, 1, 0, isQueued=1)  # 关闭吸盘
        dType.SetQueuedCmdStartExec(self.api)
        while last > dType.GetQueuedCmdCurrentIndex(self.api)[0]:
            dType.dSleep(100)
        print("物体已抓取并放置到目标位置")

    def return_init(self, velocity=50, acceleration=50):
        # dType.SetQueuedCmdClear(self.api)
        dType.SetPTPCommonParams(self.api, velocity, acceleration)
        dType.SetPTPCmd(self.api, dType.PTPMode.PTPMOVJXYZMode, 81.32083129882812, -241.886474609375, 65.85740661621094,-30, isQueued=0)
        # dType.SetQueuedCmdStartExec(self.api)
        # while last > dType.GetQueuedCmdCurrentIndex(self.api)[0]:
        #     dType.dSleep(100)

    def close(self):
        dType.DisconnectDobot(self.api)
    def get_pose(self):
        pose = dType.GetPose(self.api)
        return pose[0:4]

    def grab(self, start_pos, height=50, velocity=50, acceleration=50):
        """
        抓取物体并放置到一定高度

        参数:
            start_pos: 起始位置 [x, y, z, rHead]
            height: 抓取高度(mm)
            velocity: 速度百分比(0-100)
            acceleration: 加速度百分比(0-100)
        """
        dType.SetQueuedCmdClear(self.api)
        print(f"设置运动速度: {velocity}%，加速度: {acceleration}%")
        dType.SetPTPCommonParams(self.api, velocity, acceleration)
        # 获取起始位置
        start_x, start_y, start_z, start_r = start_pos
        target_x, target_y, target_z, target_r = start_pos
        target_z = target_z + height
        print(f"抓取物体: X={start_x}, Y={start_y}, Z={start_z}")
        dType.SetPTPCmd(self.api, dType.PTPMode.PTPMOVJXYZMode, start_x, start_y, start_z, start_r, isQueued=1)
        time.sleep(2)
        # dType.SetEndEffectorGripperEx(self.api, 1, 0, isQueued=1)  # 打开夹爪
        dType.SetEndEffectorGripperEx(self.api, 1, 1, isQueued=1)  # 关闭夹爪/吸附
        last = \
            dType.SetPTPCmd(self.api, dType.PTPMode.PTPMOVJXYZMode, target_x, target_y, target_z, target_r, isQueued=1)[
                0]
        # dType.SetEndEffectorGripperEx(self.api, 1, 0, isQueued=1)  # 打开夹爪/放下
        dType.SetQueuedCmdStartExec(self.api)
        while last > dType.GetQueuedCmdCurrentIndex(self.api)[0]:
            dType.dSleep(100)
        print("物体已抓取")


if __name__ == "__main__":
    # Example usage
    arm = ArmController()
    try:
        arm.return_init()
    finally:
        arm.close()
